Skip to content

Refactor: decouple core libraries from the default task set and make judo reusable as a third-party library - #121

Open
dta-bdai wants to merge 1 commit into
rai-opensource:mainfrom
dta-bdai:dta/refactor
Open

Refactor: decouple core libraries from the default task set and make judo reusable as a third-party library#121
dta-bdai wants to merge 1 commit into
rai-opensource:mainfrom
dta-bdai:dta/refactor

Conversation

@dta-bdai

@dta-bdai dta-bdai commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR restructures judo so its core libraries (tasks, optimizers, controllers,
simulation) are decoupled from judo's default task set and its dora application.
The goal is to let third-party packages reuse judo's mechanisms — the registries,
the Controller, the simulation/rollout backends — while registering their own
tasks, policies, and backends, without inheriting judo's batteries-included
defaults or being forced to build judo's C++ extension.

No behavior changes for the standard judo app; this is primarily a structural /
API refactor plus a set of shutdown-robustness fixes.

Key changes

Registry / application split

  • judo.app is now the batteries-included default task set. Importing it
    (or calling register_default_tasks()) registers judo's built-in tasks and
    declares how each is wired (rollout backend, simulation backend, locomotion
    policy). The individual libraries no longer auto-populate the registry, so
    third parties can register their own tasks instead of, or in addition to, the
    defaults (re-registering a name overrides it).
  • New judo.registration module with register_tasks_from_cfg,
    register_optimizers_from_cfg, and get_class_from_string for config-driven
    registration.

Registry-aware controller factory

  • New judo.controller.factory.make_controller. The Controller class is
    now decoupled from the registries — it receives its backend and locomotion
    policy explicitly. make_controller reads the registries to resolve those
    values and constructs the controller, so it depends only on the registry
    mechanism, not on any particular set of registered tasks.
  • Backends are referenced by string name at registration time; the backend
    (and its C++ extension) is only imported when a backend is actually
    instantiated. locomotion_policy_path from the registry is threaded to the
    backend via rollout_backend_kwargs["policy_path"].

Module reorganization

  • Moved the C++ extension under judo/mujoco_extensions/ (from the top-level
    mujoco_extensions/) so it ships inside the package.
  • Moved judo/app/structs.pyjudo/structs.py.
  • Renamed dora node modules for clarity: controller.pycontroller_node.py,
    simulation.pysimulation_node.py, visualization.py
    visualization_node.py.
  • Renamed hierarchical backend modules: *_simulation.py
    hierarchical_mj_simulation.py, *_rollout_backend.py
    hierarchical_mj_rollout_backend.py.

Shutdown robustness (Ctrl-C)

  • Visualizer: guard against double-removing viser handles (tab groups already
    remove their children), fixing the GuiTabHandle warning on teardown.
  • Controller node: added a _running flag and a _send_output wrapper that
    stops cleanly when the dora event stream closes mid-step.
  • CLI: main_app now catches KeyboardInterrupt / SystemExit /
    subprocess.CalledProcessError and runs a single forced cleanup.

Testing

  • pixi run -e dev python -m pytest tests -q — all tests pass.
  • pixi run -e dev ruff check judo run_mpc tests — clean.
  • pixi run build — C++ extension builds; wheel packages the .so with the
    correct RPATH.

Notes for downstream consumers

  • Applications that only want the mechanism should avoid importing judo.app
    and instead call judo.tasks.register_task (or
    judo.registration.register_tasks_from_cfg) directly.
  • Custom rollout backends registered with a locomotion_policy_path will now
    receive a policy_path kwarg — backend constructors should accept it.

@dta-bdai dta-bdai changed the title refactor Refactor: decouple core libraries from the default task set and make judo reusable as a third-party library Aug 25, 2026
@dta-bdai
dta-bdai force-pushed the dta/refactor branch 2 times, most recently from 13c2bb7 to 67810a2 Compare August 25, 2026 20:17

@bhung-bdai bhung-bdai left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The elite gold colored trace seems to have disappeared. Do you know what happened to it?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants